"""
A class to exercise the JT core functions.
"""
import itk
import vtk

import connect2
import ImageCompare


def Test_ImageNormalization():
    """
    Test image normalization in both ITK and VTK.

    From this test, it seems that only the ITK function works as expected.
    """
    import connect2
    import ImageCompare

    inputFilename='C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mha'
    outPath = 'C:/Users/bryan/bryan-code/trunk/Tests/'

    vtk2itk = connect2.vtk2itk()
    #
    # ITK Pipeline
    # ############
    input_image = itk.Image.F2.New()
    output_image = itk.Image.F2.New()

    itk_reader = itk.ImageFileReader[input_image].New()

    itk_reader.SetFileName(inputFilename)
    itk_reader.Update()

    itk_caster = itk.CastImageFilter[input_image,output_image].New()
    itk_caster.SetInput(itk_reader.GetOutput())

    itk_norm = itk.NormalizeImageFilter[output_image,output_image].New()
    itk_norm.SetInput(itk_reader.GetOutput())

    itk_SaveName = outPath + 'TestImageNormalization_itkOut.mhd'

    itk_writer = itk.ImageFileWriter[output_image].New()
    itk_writer.SetInput(itk_norm.GetOutput())
    itk_writer.SetFileName(itk_SaveName)
    itk_writer.Write()

    #
    # VTK Pipeline
    # ############
    vtk_reader = vtk.vtkMetaImageReader()
    vtk_norm = vtk.vtkImageNormalize()
    vtk_writer = vtk.vtkMetaImageWriter()
    vtk_flipper = vtk.vtkImageFlip()

    vtk_reader.SetFileName(inputFilename)
    vtk_reader.Update()

    vtk_norm.SetInput(vtk_reader.GetOutput())

    vtk_flipper.SetFilteredAxis(1)
    vtk_flipper.SetInput(vtk_norm.GetOutput())

    viewer = vtk.vtkImageViewer()
    viewer.SetInput(vtk_flipper.GetOutput())
    vtk_flipper.Update()

    vtk_SaveName = outPath + 'TestImageNormalization_vtkFlipOut.mhd'

    vtk_writer.SetInput(vtk_flipper.GetOutput())
    vtk_writer.SetFileName(vtk_SaveName)
    vtk_writer.Write()


    #Test for positive match
    try:
        #return 0 if they match, 1 otherwise.
        assert ImageCompare.compare(inputFilename,itk_SaveName)==0
        print "ITK Test Success!"
    except:
        print "ITK Test Failure!"
    #Test for positive match
    try:
        #return 0 if they match, 1 otherwise.
        assert ImageCompare.compare(inputFilename,vtk_SaveName)==0
        print "VTK Test Success!"
    except:
        print "VTK Test Failure!"



if __name__ == "__main__":
    Test_ImageNormalization()